/*  
    Payload Shuffling Technique ..

    Original Credit Goes to Cocomelonc
        => https://cocomelonc.github.io/malware/2024/09/30/malware-trick-43.html#byte-shuffling-technique
    By @5mukx
*/


use std::fs::File;
use std::io::{self, Read, Write};
use rand::{rngs::StdRng, SeedableRng, Rng};

fn shuffle_file_and_save(input_file: &str, output_file: &str, seed: u64) -> io::Result<()> {
    shuffle_bytes(input_file, output_file, seed)
}

fn main() -> io::Result<()> {
    let input_file = "w64-exec-calc-shellcode.bin";
    let shuffled_file = "shuffled.bin";
    let deshuffled_file = "deshuffled.bin";
    let seed = 12345;   

    // Lets shuffle the data and save it in seperate file ! 
    shuffle_file_and_save(input_file, shuffled_file, seed)?;

    println!("[+] shuffled and stored at [shuffled.bin]");
    // Read the shuffled data
    let mut shuffled_data = Vec::new();
    let mut shuffled_fp = File::open(shuffled_file)?;
    shuffled_fp.read_to_end(&mut shuffled_data)?;

    // At last. Deshuffle the data.
    deshuffle_bytes(&shuffled_data, shuffled_data.len(), deshuffled_file, seed)?;
    println!("[+] Desuffled and stored at [deshuffled.bin]");
    Ok(())
}


fn shuffle_bytes(input_file: &str, output_file: &str, seed: u64) -> io::Result<()> {
    // open input file for reading
    let mut in_file = File::open(input_file)?;
    let size = in_file.metadata()?.len() as usize;

    // read data from the input file
    let mut data = vec![0u8; size];
    in_file.read_exact(&mut data)?;

    // create a list of indices
    let mut indices: Vec<usize> = (0..size).collect();

    // shuffle the indices
    let mut rng = StdRng::seed_from_u64(seed);
    for i in (1..size).rev() {
        let j = rng.gen_range(0..=i);
        indices.swap(i, j);
    }

    // rearrange data based on shuffled indices
    let mut shuffled_data = vec![0u8; size];
    for (i, &idx) in indices.iter().enumerate(){
        shuffled_data[i] = data[idx];
    }

    // Write the shuffled data to the output file
    let mut out_file = File::create(output_file)?;
    out_file.write_all(&shuffled_data)?;

    Ok(())
}

fn deshuffle_bytes(data: &[u8], size: usize, output_file: &str, seed: u64) -> io::Result<()> {
    // Create a list of indices
    let mut indices: Vec<usize> = (0..size).collect();

    // Shuffle the indices to get the same order
    let mut rng = StdRng::seed_from_u64(seed);
    for i in (1..size).rev() {
        let j = rng.gen_range(0..=i);
        indices.swap(i, j);
    }

    // rearrang the data according to shuffled indices
    let mut deshuffled_data = vec![0u8; size];
    for (i, &idx) in indices.iter().enumerate() {
        deshuffled_data[idx] = data[i];
    }

    // write the deshuffled data to the output file
    let mut out_file = File::create(output_file)?;
    out_file.write_all(&deshuffled_data)?;

    Ok(())
}


